package com.llmops.core.node;

import com.llmops.core.Context;
import com.llmops.core.Node;
import com.llmops.core.NodeProperty;
import dev.langchain4j.code.CodeExecutionEngine;
import dev.langchain4j.code.graalvm.GraalVmPythonExecutionEngine;
import lombok.Data;

import java.util.List;
import java.util.Map;

public class CodeNode extends Node {
    CodeExecutionEngine engine = new GraalVmPythonExecutionEngine();
    private CodeNodeProperty property = new CodeNodeProperty();
    @Data
    public static class CodeNodeProperty extends NodeProperty {
        private String code;
        private List<String> inputs;
    }
    public CodeNode(String id, String title,Map<String, Object> config) {
        super(id,title,config);
        this.property.setCode((String)config.get("code"));
        this.property.setInputs((List<String>) config.get("inputs"));
    }

    @Override
    public void execute(Context context) {
        String call = "main(%s)";
        String param = "";
        for (int i = 0; i < property.getInputs().size(); i++) {
            if(i<property.getInputs().size()-1){
                param+="{"+property.getInputs().get(i)+"},";
            }
            param+="\"{"+property.getInputs().get(i)+"}\"";
        }
        call = String.format(call,param);
        String all = property.getCode() + System.lineSeparator() + context.fillVariable(context,call);
        String result = engine.execute(all);
        context.setVariable(getId() + "_output", result);
    }

    @Override
    public String getType() {
        return "codeNode";
    }
}